package solutions.leetcode.hard

import java.util.PriorityQueue

/**
 * @author lizhidong
 * <a href="https://leetcode.cn/problems/put-marbles-in-bags/">2551. 将珠子放入背包中</a>
 */
private class Solution2551 {
    /*
    1   3   5   1   2

    13  35  51  12
    1   1
    1       1
    1           1
        1   1
        1       1
            1   1
     */
    fun putMarbles(weights: IntArray, k: Int): Long {
        if (k == 1 || weights.size == k) {
            return 0
        }
        // 仅保留大的
        val maxQue = PriorityQueue<Long>()
        // 仅保留小的
        val minQue = PriorityQueue<Long> { o1, o2 -> if (o2 - o1 > 0) 1 else if (o2 - o1 < 0) -1 else 0 }

        // 当前最大和
        var max: Long = 0
        // 当前最小和
        var min: Long = 0
        for (i in 0 until weights.size - 1) {
            val t = weights[i] + weights[i + 1].toLong()

            if (maxQue.size < k - 1) {
                maxQue.offer(t)
                max += t
            } else {
                if (maxQue.peek() < t) {
                    max -= maxQue.poll()
                    maxQue.offer(t)
                    max += t
                }
            }
            if (minQue.size < k - 1) {
                minQue.offer(t)
                min += t
            } else {
                if (minQue.peek() > t) {
                    min -= minQue.poll()
                    minQue.offer(t)
                    min += t
                }
            }
        }
        return max - min
    }
}

fun main() {
    println(
        Solution2551().putMarbles(
            arrayOf(
                1,4,2,5,2
            ).toIntArray(), 3
        )
    )
}